import cv2
import os
import numpy as np

# 设置一个人脸的集合、id的集合
def getImagesAndLabels(path):
    # 人脸的集合
    faceSamples = []
    # id的集合
    ids = []

    # os.listdir(path) : 遍历目录下所有文件的名称
    # os.path.join(path) : 拼接“目录” + “文件名”的路径
    filePaths = [os.path.join(path,f) for f in os.listdir(path)]

    classifier = cv2.CascadeClassifier("./haarcascades/haarcascade_frontalface_alt2.xml")
    for fpath in filePaths:
        img = cv2.imread(fpath)
        gray_img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
        # 将图像中中的像素值定义为0~255 的整数
        img_numpy = np.array(gray_img,'uint8')
        faces = classifier.detectMultiScale(img_numpy,
                                    1.01,
                                    5,
                                    0,
                                    (100,100),
                                    (500,500))

        (x,y,w,h) = faces[0]
        faceSamples.append(img_numpy[y:y+h,x:x+w])
        # os.path.split 将 ./train/xxxx.jpg => train + xxx.jpg
        # split(".")[0] 将 aaa.bbb.jpg => aaa + bbb + jpg
        id = int(os.path.split(fpath)[1].split(".")[0])
        ids.append(id)
    return faceSamples,ids

if __name__ == '__main__':
    faceSamples,ids = getImagesAndLabels("./train/")
    # 训练人脸的识别器
    recognizer = cv2.face.LBPHFaceRecognizer_create()
    # 训练人脸数据以及对应的id
    recognizer.train(faceSamples,np.array(ids))
    # 保存人脸的训练结果
    recognizer.write('./data/train.yml')
